package org.apache.flink.runtime.resourceprovider;

import java.util.Collection;
import java.util.Optional;

import org.apache.flink.api.common.ConsumerID;
import org.apache.flink.runtime.clusterframework.types.ResourceID;
import org.apache.flink.runtime.providermanager.ProviderManagerActions;
import org.apache.flink.runtime.resourceconsumer.ResourceConsumerGateway;
import org.apache.flink.runtime.resourceconsumer.ConsumingMasterId;
import org.apache.flink.util.function.SupplierWithException;

/**
 * A {@link ConsumerTable ConsumerTable’s} task is to manage the lifecycle of a consumer on the {@link
 * ResourceProvider}.
 *
 * <p>There can always only be at most one consuming per {@link ConsumerID}. In order to create a {@link Consumer}
 * one needs to provide a {@link ConsumerTable.ConsumerServices} instance which is owned by the job.
 *
 * <p>A consuming can be connected to a leading ConsumerManager or can be disconnected. In order to establish a
 * connection, one needs to call {@link Consumer#connect}. Once a consuming is connected, the respective {@link
 * Connection} can be retrieved via its {@link ConsumerID} or via the {@link ResourceID} of the leader. A
 * connection can be disconnected via {@link Connection#disconnect}.
 *
 * <p>In order to clean up a {@link Consumer} one first needs to disconnect from the leading JobManager.
 * In order to completely remove the {@link Consumer} from the {@link ConsumerTable}, one needs to call {@link
 * Consumer#close} which also closes the associated {@link ConsumerTable.ConsumerServices} instance.
 */
public interface ConsumerTable extends AutoCloseable {

    /**
     * Gets a registered {@link Consumer} or creates one if not present.
     *
     * @param consumerId consumerId identifies the consumer to get
     * @param consumerServicesSupplier jobServicesSupplier create new {@link ConsumerTable.ConsumerServices} if a
     *     new job needs to be created
     * @return the current job (existing or created) registered under jobId
     * @throws E if the job services could not be created
     */
    <E extends Exception> Consumer getOrCreateConsumer(
            ConsumerID consumerId,
            SupplierWithException<? extends ConsumerTable.ConsumerServices, E> consumerServicesSupplier)
            throws E;

    /**
     * Gets the job registered under jobId.
     *
     * @param consumerId jobId identifying the job to get
     * @return an {@code Optional} containing the {@link Consumer} registered under jobId, or an empty
     *     {@code Optional} if no job has been registered
     */
    Optional<Consumer> getConsumer(ConsumerID consumerId);

    /**
     * Gets the connection registered under jobId.
     *
     * @param jobId jobId identifying the connection to get
     * @return an {@code Optional} containing the {@link Connection} registered under jobId, or an
     *     empty {@code Optional} if no connection has been registered (this could also mean that a
     *     job which has not been connected exists)
     */
    Optional<Connection> getConnection(ConsumerID consumerId);

    /**
     * Gets the connection registered under resourceId.
     *
     * @param resourceId resourceId identifying the connection to get
     * @return an {@code Optional} containing the {@link Connection} registered under resourceId, or
     *     an empty {@code Optional} if no connection has been registered
     */
    Optional<Connection> getConnection(ResourceID resourceId);

    /**
     * Gets all registered jobs.
     *
     * @return collection of registered jobs
     */
    Collection<Consumer> getConsumers();

    /**
     * Returns {@code true} if the job table does not contain any jobs.
     *
     * @return {@code true} if the job table does not contain any jobs, otherwise {@code false}
     */
    boolean isEmpty();

    /**
     * A job contains services which are bound to the lifetime of a resource consumer. Moreover, it can be
     * connected to a leading ResourceConsumer and then store further services which are bound to the
     * lifetime of the ResourceConsumer connection.
     *
     * <p>Accessing any methods after a job has been closed will throw an {@link
     * IllegalStateException}.
     */
    interface Consumer {

        /**
         * Returns {@code true} if the job is connected to a JobManager.
         *
         * @return {@code true} if the job is connected to a JobManager, otherwise {@code false}
         */
        boolean isConnected();

        /**
         * Returns the {@link ConsumerID} which is associated with this consumer.
         *
         * @return job id which is associated with this job
         */
        ConsumerID getConsumerId();

        /**
         * Returns the associated {@link Connection} if the job is connected to a JobManager.
         *
         * @return an {@code Optional} containing the associated {@link Connection} instance if the
         *     job is connected to a leading JobManager, or an empty {@code Optional} if the job is
         *     not connected
         */
        Optional<Connection> asConnection();

        /**
         * Connects to a ResouceConsumer and associates the provided services with this
         * connection.
         *
         * <p>A job can only be connected iff {@code Job#isConnected() == false}.
         *
         * @param resourceId resourceId of the JobManager to connect to
         * @param resourceConsumerGateway jobMasterGateway of the JobManager to connect to
         * @param taskManagerActions taskManagerActions associated with this connection
         * @param checkpointResponder checkpointResponder associated with this connection
         * @param aggregateManager aggregateManager associated with this connection
         * @param partitionStateChecker partitionStateChecker associated with this connection
         * @return the established {@link Connection}
         * @throws IllegalStateException if the job is already connected
         */
        Connection connect(
                ResourceID resourceId,
                ResourceConsumerGateway resourceConsumerGateway,
                ProviderManagerActions taskManagerActions);

        /** Closes this job and removes it from the owning {@link ConsumerTable}. */
        void close();
    }

    /**
     * A connection contains services bound to the lifetime of a connection with a JobManager.
     *
     * <p>A connection can be disconnected by calling {@link #disconnect()}. Disconnecting a
     * connection will close all services bound to the connection and return the remaining job
     * instance.
     *
     * <p>Accessing any methods after a connection has been disconnected will throw an {@link
     * IllegalStateException}.
     */
    interface Connection {

        /**
         * Disconnects the connection, closing all associated services thereby and returning the
         * remaining job.
         *
         * @return the remaining job belonging to this connection
         */
        Consumer disconnect();

        ConsumingMasterId getConsumingMasterId();

        ResourceConsumerGateway getResourceConsumerGateway();

        ProviderManagerActions getTaskManagerActions();

        ConsumerID getConsumerId();

        ResourceID getResourceId();
    }

    /**
     * Services associated with a consumer. The services need to provide a {@link
     * LibraryCacheManager.ClassLoaderHandle} and will be closed once the associated {@link
     * ConsumerTable.Consumer} is being closed.
     */
    interface ConsumerServices {

        /**
         * Closes the job services.
         *
         * <p>This method is called once the {@link ConsumerTable.Consumer} is being closed.
         */
        void close();
    }
}
